#include <bits/stdc++.h>
#include <windows.h>

using namespace std;

bool readBMP(const char *filename, vector<unsigned char> &grayPixels, int &width, int &height)
{
    ifstream file(filename, ios::binary);
    if (!file)
    {
        return false;
    }
    BITMAPFILEHEADER fileHeader;
    BITMAPINFOHEADER infoHeader;

    file.read(reinterpret_cast<char *>(&fileHeader), sizeof(fileHeader));
    if (fileHeader.bfType != 0x4D42)
        return false; // 'BM'

    file.read(reinterpret_cast<char *>(&infoHeader), sizeof(infoHeader));
    if (infoHeader.biCompression != 0)
        return false; // 仅支持无压缩

    width = infoHeader.biWidth;
    height = abs(infoHeader.biHeight);
    bool isTopDown = infoHeader.biHeight < 0;
    if (isTopDown)
        height = -infoHeader.biHeight;

    int bitCount = infoHeader.biBitCount;
    if (bitCount != 8 && bitCount != 24)
        return false; // 仅支持8位和24位

    // 读取调色板（仅8位）
    vector<unsigned char> palette;
    if (bitCount == 8)
    {
        palette.resize(1024);
        file.read(reinterpret_cast<char *>(palette.data()), 1024);
    }

    // 计算行大小（包括填充）
    int rowSize = (((width * bitCount + 31) / 32) * 4);
    vector<unsigned char> row(rowSize);

    grayPixels.resize(width * height);

    // 读取像素数据
    for (int y = 0; y < height; ++y)
    {
        int actualY = isTopDown ? y : height - 1 - y;
        file.read(reinterpret_cast<char *>(row.data()), rowSize);

        if (bitCount == 8)
        {
            for (int x = 0; x < width; ++x)
            {
                unsigned char index = row[x];
                unsigned char B = palette[index * 4];
                unsigned char G = palette[index * 4 + 1];
                unsigned char R = palette[index * 4 + 2];
                grayPixels[actualY * width + x] = static_cast<unsigned char>(0.299 * R + 0.587 * G + 0.114 * B + 0.5);
            }
        }
        else if (bitCount == 24)
        {
            for (int x = 0; x < width; ++x)
            {
                int offset = x * 3;
                unsigned char B = row[offset];
                unsigned char G = row[offset + 1];
                unsigned char R = row[offset + 2];
                grayPixels[actualY * width + x] = static_cast<unsigned char>(0.299 * R + 0.587 * G + 0.114 * B + 0.5);
            }
        }
    }

    return true;
}

bool writeBMP(const char *filename, const vector<unsigned char> &grayPixels, int width, int height)
{
    ofstream file(filename, ios::binary);
    if (!file)
        return false;

    int rowSize = ((width * 24 + 31) / 32) * 4;
    int imageSize = rowSize * height;

    BITMAPFILEHEADER fileHeader = {};
    fileHeader.bfType = 0x4D42;
    fileHeader.bfSize = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER) + imageSize;
    fileHeader.bfOffBits = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER);

    BITMAPINFOHEADER infoHeader = {};
    infoHeader.biSize = sizeof(BITMAPINFOHEADER);
    infoHeader.biWidth = width;
    infoHeader.biHeight = height;
    infoHeader.biPlanes = 1;
    infoHeader.biBitCount = 24;
    infoHeader.biSizeImage = imageSize;

    file.write(reinterpret_cast<const char *>(&fileHeader), sizeof(fileHeader));
    file.write(reinterpret_cast<const char *>(&infoHeader), sizeof(infoHeader));

    vector<unsigned char> row(rowSize, 0);
    for (int y = 0; y < height; ++y)
    {
        int srcY = height - 1 - y;
        for (int x = 0; x < width; ++x)
        {
            unsigned char gray = grayPixels[srcY * width + x];
            int offset = x * 3;
            row[offset] = gray;     // B
            row[offset + 1] = gray; // G
            row[offset + 2] = gray; // R
        }
        file.write(reinterpret_cast<const char *>(row.data()), rowSize);
    }

    return true;
}

void applyLinearTransform(vector<unsigned char> &pixels, double a, double b)
{
    for (auto &p : pixels)
    {
        double newVal = a + (b - a) / 255 * p;
        newVal = max(0.0, min(255.0, newVal));
        p = static_cast<unsigned char>(newVal + 0.5);
    }
}

void applyLogTransform(vector<unsigned char> &pixels, double a, double b)
{
    for (auto &p : pixels)
    {
        double newVal = a + log(p + 1.0) / b;
        newVal = max(0.0, min(255.0, newVal));
        p = static_cast<unsigned char>(newVal + 0.5);
    }
}

void applyHistogramEqualization(vector<unsigned char> &pixels, int width, int height)
{
    vector<int> hist(256, 0);
    for (auto p : pixels)
        hist[p]++;

    vector<int> cdf(256);
    cdf[0] = hist[0];
    for (int i = 1; i < 256; ++i)
        cdf[i] = cdf[i - 1] + hist[i];

    int cdf_min = *find_if(cdf.begin(), cdf.end(), [](int x)
                           { return x > 0; });
    int total = width * height;

    vector<unsigned char> lut(256);
    for (int i = 0; i < 256; ++i)
    {
        if (cdf[i] == 0)
            lut[i] = 0;
        else
        {
            double val = 255.0 * (cdf[i] - cdf_min) / (total - cdf_min);
            lut[i] = static_cast<unsigned char>(val + 0.5);
        }
    }

    for (auto &p : pixels)
        p = lut[p];
}

void applySmoothing(vector<unsigned char> &pixels, int width, int height, int windowSize)
{
    vector<unsigned char> result(pixels.size());
    int radius = windowSize / 2;

    for (int y = 0; y < height; ++y)
    {
        for (int x = 0; x < width; ++x)
        {
            int sum = 0, count = 0;
            for (int dy = -radius; dy <= radius; ++dy)
            {
                int ny = y + dy;
                if (ny < 0 || ny >= height)
                    continue;
                for (int dx = -radius; dx <= radius; ++dx)
                {
                    int nx = x + dx;
                    if (nx < 0 || nx >= width)
                        continue;
                    sum += pixels[ny * width + nx];
                    count++;
                }
            }
            result[y * width + x] = count ? static_cast<unsigned char>(sum / count) : pixels[y * width + x];
        }
    }

    pixels.swap(result);
}

int main(int argc, char *argv[])
{
    if (argc < 4)
    {
        cerr << "Usage: imgen [options] <input.bmp> <output.bmp>\n";
        return 1;
    }

    string option = argv[1];

    if (option == "-g")
    {
        if (argc != 7)
            throw invalid_argument("Invalid arguments for -g");
        string transform = argv[2];
        double a = stod(argv[3]), b = stod(argv[4]);
        vector<unsigned char> pixels;
        int width, height;
        if (!readBMP(argv[5], pixels, width, height))
            throw runtime_error("Failed to read BMP");
        if (transform == "-l")
            applyLinearTransform(pixels, a, b);
        else if (transform == "-e")
            applyLogTransform(pixels, a, b);
        else
            throw invalid_argument("Invalid transform type");
        if (!writeBMP(argv[6], pixels, width, height))
            throw runtime_error("Failed to write BMP");
    }
    else if (option == "-h")
    {
        if (argc != 4)
            throw invalid_argument("Invalid arguments for -h");
        vector<unsigned char> pixels;
        int width, height;
        if (!readBMP(argv[2], pixels, width, height))
            throw runtime_error("Failed to read BMP");
        applyHistogramEqualization(pixels, width, height);
        if (!writeBMP(argv[3], pixels, width, height))
            throw runtime_error("Failed to write BMP");
    }
    else if (option == "-s")
    {
        if (argc != 5)
            throw invalid_argument("Invalid arguments for -s");
        int windowSize = (string(argv[3]) == "-3") ? 3 : 5;
        vector<unsigned char> pixels;
        int width, height;
        if (!readBMP(argv[2], pixels, width, height))
            throw runtime_error("Failed to read BMP");
        applySmoothing(pixels, width, height, windowSize);
        if (!writeBMP(argv[4], pixels, width, height))
            throw runtime_error("Failed to write BMP");
    }
    else
    {
        throw invalid_argument("Unknown option");
    }
    cout << "complete" << endl;

    return 0;
}